向外借力：Pluto助力MLIR编译器的多面体优化

摘要

多面体编译是一项成熟的编译优化技术，演进了几十年，在传统的编译器中常作为一种优化工具使用，比如LLVM中使用的Polly，在GCC中使用的GRAPHITE。近些年来，多面体技术也引入到AI编译器中，进行循环优化及算子融合优化等。本文将关注在MLIR中以类插件的形式引入多面体优化技术，补充其多面体优化能力。

多面体模型的介绍

多面体模型(Polyhedral)主要关注的是程序设计中的循环优化问题，两层循环的循环变量的取值范围可以构成一个平面，三层循环的循环变量可以组成一个长方体，如图1所示，因此得名多面体模型。

图1 不同循环层次的多面体表示[1]

多面体编译优化关注的是在确保程序执行正确的前提下重组多重循环的结构，实现性能的最优化。比如图2的循环中，左图表示的原始的二维迭代空间，蓝色箭头表示数据（黑点）之间存在依赖关系，对角线的绿色表示数据没有依赖关系，经过变基操作之后变为右图的表达式及迭代空间，从形状看像是把多面体进行了变形，形象地体现出多面体优化的过程。当然，变形的目的是为了实现并行计算，达到更好的性能，具体分析可以参考[1][2]。

图2 多面体的变基转换及Affine_map表示[2]

MLIR中的多面体表示

MLIR关于多面体的设计重表达轻优化，也就是说MLIR充分利用其IR的特性定义了表示多面体的方言Affine[3]，而没有进行多面体的优化实现。这样做的目的符合MLIR重在编译器基础设施搭建的特性，而具体的实现可以自行定义，好比MLIR向用户交付了布局合理的毛坯房，内部装修各取所好。另外，关于多面体优化的工具也很成熟，各种开源工具齐全，比如isl[4]，Polly[5]，Pluto[6]，以及CLooG[7]，也为将编译优化工具引入到MLIR提供了便利。

MLIR中Affine方言的定义使用具有多面体特征的循环和条件判断来表示，显示地表示静态控制部分（SCoP：Static Control Part），比如affine.for, affine.if, affine.parallel等，具体可以参见官方文档[3]。在Affine的表达中，使用Dimension和Symbol两类标识符，二者在MLIR语法中均为index类型，同时MLIR也对这两类表示进行约束有助于提升分析和转换能力。从表示形式上看，Dimension以圆括号来声明，Symbol用方括号声明，Dimension即字面意思表示仿射对象的维度信息，比如映射，集合或者具体的loop循环以及一个tensor，Symbol表示的是多面体中的参数，在编译阶段是未知的。在准线性的分析表达上，Symbol当作常量对待，因此Symbol和Dimension之间可以进行乘加等线性操作，但Dimension之间的操作是非法的。另外，Dimension和symbol都遵从SSA赋值。

MLIR的Affine重表达体现在通过具体的映射可以表示出多面体的变换，比如图1中的变基操作，通过affine_map语句就能够体现出来。矩阵计算中常用的关于内存的tiling操作，也可以通过affine_map表示。通过#tiled_2d_128x256 = affine_map<(d0, d1) -> (d0 div 128, d1 div 256, d0 mod 128, d1 mod 256)>可以表示如下图所示的tiling切分。X维度按照128个元素，Y维度按照256维度进行tiling切分，形成了x.outer和y.outer，内部的取模运算计算在单个tiling内部的偏移量，形成x.inner和y.inner。图4中表示采用Affine方言表示多项式乘法C[i+j] += A[i] * B[j]，从中可以明确看到和多面体相关的操作表达。

图3 tiling的操作

图4 多项式乘法的Affine方言表示

MLIR中引入多面体优化

Polygeist[9]文章沿用LLVM引入Polly，GCC引入GRAPHITE的思路，将开源的Pluto引入到MLIR中实现多面体优化。整体的编译采用LLVM的编译流程，前端Clang分析输入的C语言代码，转换到MLIR中的Affine方言，然后从Affine转换到Pluto，二者之间的交互采用多面体技术中常用的数据格式OpenScop[8]，优化后的代码再次转换到MLIR，然后走LLVM的编译流程。

在整体的转换中，使用的方法就是遍历语法树（AST）。Clang前端遍历Clang的AST语法树，将node映射到MLIR的操作中；从Pluto到MLIR的转换过程中使用CLooG生成初始循环的AST，然后遍历AST中的node，创建与循环和条件判断相对应的MLIR的操作。这种做法可以产生和现有编译器的应用程序二进制接口（Application Binary Interface）兼容的代码，免去重新构建基础设施。

Polygeist出彩的地方的是将C/C++引入到了MLIR，实现了clang前端对MLIR的对接，利用MLIR的IR变换能力对接到开源的Pluto工具进行多面体的优化。从整体的测试效果看，优化的性能没有得到特别显著的提升，原因在于针对benchmark中的实例，现有的编译工具已经优化的很好，比如LLVM对于Clang前端的IR能够减少特定的load操作，而LLVM对MLIR的IR还不支持此优化功能，MLIR还是相对比较新鲜的事物。正如文中所言，使用MLIR生成比多面体工具更加优化的代码不是该工作的兴趣点，而是证实在MLIR中引入多面体优化工具的灵活性。

当然，整体的转换工作不是一蹴而就的，需要涉及到C语言和MLIR之间类型的转换，目前支持的数据类型如图5所示，不支持struct的数据类型，同时也不支持C语言中的break操作。在MLIR中没有指针的概念，关于内存的表示只有memref，MLIR本身也不支持在memref中嵌套memref，为了对接C语言中的指针的指针，文章增加对memref的嵌套使用，也就是修改了原有MLIR的特性。

图5 C语言，LLVM IR和MLIR数据类型的对应关系

整体实现的流程如图6所示，前端搭建了C代码到MLIR的连接，通过遍历Clang的AST语法树，将每个访问的节点映射到MLIR中的SCF或者standard方言中的操作。MLIR起到表达控制流的作用，在方言的表达中直接查找到循环，减少在Pluto的CFG (Control Flow Graph) 中查找loops循环的必要。但是实际上，Pluto还是参与了C/C++中非线性for循环的查找。AST中C语言的数据类型先是转换到LLVM的数据类型，然后转换到MLIR中Standard方言中的数据类型。

多面体相关的Affine code转换通过识别标识符将#pragma scop和#pragma endscop表示的code直接转化为affine.for 循环。循环约束以affine_map的形式表示，比如(affine_map<()[s0]->(s0)>[%bound])，()中表示的是表示Dimension，[ ]表示的是Symbol，如前文所言。

前端输入的C语言经过转换到MLIR-SCF方言层级，通过raise-affine的PASS转换为Affine方言，具体实现的功能是将standard 方言中的load, store，SCF方言中的for和if转换到affine方言中对应的操作。利用Pluto的能力进行优化处理，然后再通过low-affine PASS转回到MLIR-SCF，此处借助CLooG进行语法树分析，然后走MLIR的LLVM编译流程。

图6 Polygeist的编译流程

结论与思考

现有的多面体优化的库，比如isl，Polly，主要用于C语言的source-source转换，聚焦于底层级的优化，无法直接用在MLIR的设计中，因为底层级的表示无法还原完整的高层级的表达。Polygeist的工作将MLIR的多面体表示和现有的高层级的优化工具结合起来，采用MLIR Affine方言和OpenScop数据格式的双向转换方便开发者搭建基于MLIR的编译流程，然后使用现有的多面体优化工具优化，最后返回到MLIR进行进一步的转换并最终生成代码。对于我们的启发在于可以在MLIR中引入其他优化工具助力编译优化，根据需求补足MLIR中缺失的能力。

多面体优化是一项成熟的技术，但也受限于对仿射变换的依赖，对无法进行仿射的循环的优化能力较弱，存在一定的局限性，因此无法在工业界得到广泛应用。同时，多面体优化技术理论相对复杂难懂，从事相关研究的人员较少，难以进行落地。尽管如此，多面体技术在解决特定的问题方面尤其独特的作用，比如在深度学习领域，对多算子融合和多层循环优化方面有着极大的帮助，可以将现有的多面体技术引入到AI编译器中，进行特定功能的优化。

由于水平有限，文中存在不足的地方请各位读者批评指正，也欢迎大家一起参与我们的讨论。

参考文献

[1] Austin Derrow-Pinion, Jennifer She, David Wong, et al. ETA Predictionwith Graph Neural Networks in Google Maps. 2021

[1] Uday Bondhugula, Polyhedral Compilation Opportunities in MLIR http://impact.gforge.inria.fr/impact2020/slides/IMPACT_2020_keynote.pdf

[2] 要术甲杰, Polyhedral Model—AI芯片软硬件优化利器https://mp.weixin.qq.com/s?__biz=MzI3MDQ2MjA3OA==&mid=2247485130&idx=1&sn=a5773bf17e6854d1238b035366641bcc&chksm=ead1fbdbdda672cdf9b2480a431cef85e4d377d07f8c586a932adabd50656cbdcd7d891156bf&mpshare=1&scene=1&srcid=&sharer_sharetime=1569677798809&sharer_shareid=b33ef36fa0caf5cb82e76916516aa7df#rd

[3] https://mlir.llvm.org/docs/Dialects/Affine/#affine-maps

[4] Sven Verdoolaege. 2010. isl: An integer set library for the polyhedral model. In International Congress on Mathematical Software. Springer, 299–302

[5] Tobias Grosser, Armin Groesslinger, and Christian Lengauer. 2012. Polly—performing polyhedral optimizations on a low-level intermediate representation. Parallel Processing Letters 22, 04 (2012), 1250010.

[6] Uday Bondhugula, Albert Hartono, Jagannathan Ramanujam, and Ponnuswamy Sadayappan. 2008. A Practical Automatic Polyhedral Parallelizer and Locality Optimizer. ACM SIGPLAN Notices 43, 6 (2008), 101–113.

[7] Uday Bondhugula, Albert Hartono, Jagannathan Ramanujam, and Ponnuswamy Sadayappan. 2008. A Practical Automatic Polyhedral Parallelizer and Locality Optimizer. ACM SIGPLAN Notices 43, 6 (2008),

101–113.

[8] Cédric Bastoul. 2011. Openscop: A specification and a library for data exchange in polyhedral compilation tools. Technical Report. Paris-Sud University.

 往期推荐

关于壁仞科技研究院

壁仞科技研究院作为壁仞科技的前沿研究部门，旨在研究新型智能计算系统的关键技术，重点关注新型架构，先进编译技术和设计方法学，并将逐渐拓展研究方向，探索未来智能系统的各种可能。壁仞科技研究院秉持开放的原则，将积极投入各类产学研合作并参与开源社区的建设，为相关领域的技术进步做出自己的贡献。

扫码关注我们